package com.zhu.oj.service.impl;

import com.zhu.oj.service.AIService;
import com.zhu.oj.service.adapter.AiAdapter;
import io.github.pigmesh.ai.deepseek.core.chat.ChatCompletionResponse;
import org.springframework.stereotype.Service;

import com.zhu.oj.common.constant.CommonConstant;
import com.zhu.oj.model.vo.AIChatStreamResponse;
import io.github.pigmesh.ai.deepseek.config.DeepSeekProperties;
import io.github.pigmesh.ai.deepseek.core.DeepSeekClient;
import io.github.pigmesh.ai.deepseek.core.Json;
import io.github.pigmesh.ai.deepseek.core.chat.ChatCompletionRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.function.Function;

@Service
@Slf4j
public class AIServiceImpl implements AIService {

    @Autowired
    private DeepSeekClient deepSeekClient;

    @Autowired
    private DeepSeekProperties deepSeekProperties;

    Function<String, String> elt = s -> s.replaceAll("<think>[\\s\\S]*?</think>", "").replaceAll("\n", "");


    @Override
    public Flux<AIChatStreamResponse> getAiChatStreamResponseFlux(String prompt, String cycelCount, Map<String, String> userContactContext) {
        ChatCompletionRequest request = ChatCompletionRequest.builder().model(deepSeekProperties.getModel())
                .addUserMessage(prompt)
                .addAssistantMessage(elt.apply(userContactContext.getOrDefault(cycelCount, "")))
                .addSystemMessage(CommonConstant.AI_SYSTEM_PROMPT).maxCompletionTokens(5000).build();
        log.info("request {}", Json.toJson(request));
        userContactContext.remove(cycelCount);

        Flux<ChatCompletionResponse> originalFlux = deepSeekClient.chatFluxCompletion(request)
                .doOnNext(response -> {
                    String contentChunk = AiAdapter.choicesProcess.apply(response.choices());
                    if (contentChunk != null && !contentChunk.isEmpty()) {
                        userContactContext.merge(cycelCount, contentChunk, String::concat);
                    }
                });

        return AiAdapter.buildAIChatStreamResponse(originalFlux);
    }

    @Override
    public Flux<AIChatStreamResponse> questionTips(String prompt) {
        prompt += CommonConstant.QUESTION_TIPS_PROMPTS;
        Flux<ChatCompletionResponse> chatCompletionResponseFlux = deepSeekClient.chatFluxCompletion(prompt);
        return AiAdapter.buildAIChatStreamResponse(chatCompletionResponseFlux);
    }
}
